import cv2
import os
from cv2 import mean
import numpy as np
import sys
sys.path.append(os.path.abspath('.'))
from utils.util import sci_SSIM as SSIM 
from utils.util import SSIM_mean_param,SSIM_std_param,SSIM_cov_param
from .base_util import BaseUtil
from scipy import signal

class PolarTianyuUtil(BaseUtil):
    def __init__(self, opt):
        BaseUtil.__init__(self,opt)
        
    def phaseData(self):
        record = []
        self.img = {}
        self.NV_dir = "/".join(self.opt.dir.split("/")[:-2])
        if not os.path.exists(self.opt.dir):
            return False
        self.cklogfie_dir = os.path.join("/".join(self.opt.dir.replace("results", "datasets/checkpoints").split("/")[:7]),"loss_log.txt")
        if not os.path.exists(self.cklogfie_dir):
            return False
        
        return True
    
    def process(self):
        self.logfile = os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'uilt_log.txt')
        if os.path.exists(self.logfile):
            os.remove(self.logfile)
        for process in self.opt.processes:
            getattr(self,process)()
        # self.phaseData()
        
    def train_val_curve(self):
        abstract_str = lambda x,a,b:x[x.find(a)+len(a):][:x[x.find(a)+len(a):].find(b)]
        if "SD" in self.cklogfie_dir:
            self.SDcklogfile_dir = self.cklogfie_dir
            self.S0cklogfile_dir = self.cklogfie_dir.replace("SD", "S0")
        elif "S0" in self.cklogfie_dir:
            self.S0cklogfile_dir = self.cklogfie_dir
            self.SDcklogfile_dir = self.cklogfie_dir.replace("S0", "SD")
        with open(self.SDcklogfile_dir,'r') as f:
            x = f.read()
            xlines = x.split("\n")
            SD_dict = {}
            SD_dict["Training"] = {}
            SD_dict["Validating"] = {}
            mode = "Training"
            parameters = ["Accurate","F_score","APCER","NPCER","ACER"]
            funcs = ["max","max","min","min","min"]
            SD_max_dict = {}
            SD_max_dict["Training"] = {}
            SD_max_dict["Validating"] = {}
            for xline in xlines:
                if "================ Training Loss" in xline:
                    SD_dict = {}
                    SD_dict["Training"] = {}
                    SD_dict["Validating"] = {}
                if '----------Training----------------' in xline:
                    mode = "Training"
                    pass
                elif '----------Validating----------------' in xline:
                    mode = "Validating"
                    pass
                for (param,func) in zip(parameters,funcs):
                    if param + ":" in xline:
                        if param not in SD_dict[mode].keys():
                            SD_dict[mode][param] = []
                            SD_max_dict[mode][param] = []
                            SD_max_dict[mode][param].append(float(abstract_str(xline,param + ": "," ")))
                        else:
                            SD_max_dict[mode][param].append(eval(func)(eval(func)(SD_max_dict[mode][param]),float(abstract_str(xline,param + ": "," "))))
                        SD_dict[mode][param].append(float(abstract_str(xline,param + ": "," ")))
                        
        with open(self.S0cklogfile_dir,'r') as f:
            x = f.read()
            xlines = x.split("\n")
            mode = "Training"
            parameters = ["Accurate","F_score","APCER","NPCER","ACER"]
            funcs = ["max","max","min","min","min"]
            for xline in xlines:
                if "================ Training Loss" in xline:
                    S0_dict = {}
                    S0_dict["Training"] = {}
                    S0_dict["Validating"] = {}
                    S0_max_dict = {}
                    S0_max_dict["Training"] = {}
                    S0_max_dict["Validating"] = {}
                elif '----------Training----------------' in xline:
                    mode = "Training"
                    pass
                elif '----------Validating----------------' in xline:
                    mode = "Validating"
                    pass
                for (param,func) in zip(parameters,funcs):
                    if param + ":" in xline:
                        # if param not in S0_dict[mode].keys():
                        #     S0_dict[mode][param] = []
                        if param not in S0_dict[mode].keys():
                            S0_dict[mode][param] = []
                            S0_max_dict[mode][param] = []
                            S0_max_dict[mode][param].append(float(abstract_str(xline,param + ": "," ")))
                        else:
                            S0_max_dict[mode][param].append(eval(func)(eval(func)(S0_max_dict[mode][param]),float(abstract_str(xline,param + ": "," "))))
                        S0_dict[mode][param].append(float(abstract_str(xline,param + ": "," ")))
                        # S0_dict[mode][param].append(float(abstract_str(xline,param + ": "," ")))
                        # else:
                        #     S0_dict[mode][param].append(eval(func)(eval(func)(S0_dict[mode][param]),float(abstract_str(xline,param + ": "," "))))
        
        import matplotlib.pyplot as plt
        
        # counts = range(len(ans_dict[mode][param]))
        # mode = "Training"
        colors = ['r','b']
        linestyles = ['-','--']
        # parameters = ["F_score","APCER","NPCER","ACER"]
        for param in parameters:
            plt.cla()
            wd=31
            k=4
            ind1 = len(SD_dict["Training"][param])
            ind2 = len(S0_dict["Training"][param])
            ind3 = len(SD_dict["Validating"][param])
            ind4 = len(S0_dict["Validating"][param])
            print(ind1,ind2,ind3,ind4)
            epoch1 = max(ind1,ind2)
            epoch2 = max(ind3,ind4)

            y = self.checkepoch(SD_dict["Training"][param],epoch1,wd,k)
            plt.plot(range(y.shape[0]),y,'r-')
            # y = signal.savgol_filter(SD_dict["Validating"][param],wd,k)
            y = self.checkepoch(SD_dict["Validating"][param],epoch2,wd,k)
            plt.plot(range(y.shape[0]),y,'r--')
            # y = signal.savgol_filter(S0_dict["Training"][param],wd,k)
            y = self.checkepoch(S0_dict["Training"][param],epoch1,wd,k)
            plt.plot(range(y.shape[0]),y,'b-')
            # y = signal.savgol_filter(S0_dict["Validating"][param],wd,k)
            y = self.checkepoch(S0_dict["Validating"][param],epoch2,wd,k)
            plt.plot(range(y.shape[0]),y,'b--')
            plt.xlabel("epoch")
            plt.legend(["Polarized Training","Polarized Testing","Visible Training", "Visible Testing"])
            plt.savefig(os.path.join(self.opt.dir,self.opt.name + "_v" + self.opt.version + "_" +param+"_TrueCurve.jpg"))
        
        for param in parameters:
            plt.cla()
            wd=31
            k=4
            ind1 = len(SD_dict["Training"][param])
            ind2 = len(S0_dict["Training"][param])
            ind3 = len(SD_dict["Validating"][param])
            ind4 = len(S0_dict["Validating"][param])
            print(ind1,ind2,ind3,ind4)
            epoch1 = max(ind1,ind2)
            epoch2 = max(ind3,ind4)

            y = self.checkepoch(SD_dict["Training"][param],epoch1,wd,k,True)
            plt.plot(range(y.shape[0]),y,'r-')
            # y = signal.savgol_filter(SD_dict["Validating"][param],wd,k)
            y = self.checkepoch(SD_dict["Validating"][param],epoch2,wd,k,True)
            plt.plot(range(y.shape[0]),y,'r--')
            # y = signal.savgol_filter(S0_dict["Training"][param],wd,k)
            y = self.checkepoch(S0_dict["Training"][param],epoch1,wd,k,True)
            plt.plot(range(y.shape[0]),y,'b-')
            # y = signal.savgol_filter(S0_dict["Validating"][param],wd,k)
            y = self.checkepoch(S0_dict["Validating"][param],epoch2,wd,k,True)
            plt.plot(range(y.shape[0]),y,'b--')
            plt.xlabel("epoch")
            plt.legend(["Polarized Training","Polarized Testing","Visible Training", "Visible Testing"])
            plt.savefig(os.path.join(self.opt.dir,self.opt.name + "_v" + self.opt.version + "_" +param+"_curve.jpg"))
        
        for param in parameters:
            plt.cla()
            wd=31
            k=4
            ind1 = len(SD_max_dict["Training"][param])
            ind2 = len(S0_max_dict["Training"][param])
            ind3 = len(SD_max_dict["Validating"][param])
            ind4 = len(S0_max_dict["Validating"][param])
            print(ind1,ind2,ind3,ind4)
            epoch1 = max(ind1,ind2)
            epoch2 = max(ind3,ind4)

            y = self.checkepoch(SD_max_dict["Training"][param],epoch1,wd,k)
            plt.plot(range(y.shape[0]),y,'r-')
            # y = signal.savgol_filter(SD_dict["Validating"][param],wd,k)
            y = self.checkepoch(SD_max_dict["Validating"][param],epoch2,wd,k)
            plt.plot(range(y.shape[0]),y,'r--')
            # y = signal.savgol_filter(S0_dict["Training"][param],wd,k)
            y = self.checkepoch(S0_max_dict["Training"][param],epoch1,wd,k)
            plt.plot(range(y.shape[0]),y,'b-')
            # y = signal.savgol_filter(S0_dict["Validating"][param],wd,k)
            y = self.checkepoch(S0_max_dict["Validating"][param],epoch2,wd,k)
            plt.plot(range(y.shape[0]),y,'b--')
            plt.xlabel("epoch")
            plt.legend(["Polarized Training","Polarized Testing","Visible Training", "Visible Testing"])
            plt.savefig(os.path.join(self.opt.dir,self.opt.name + "_v" + self.opt.version + "_" +param+"_BestCurve.jpg"))
        pass

    def mean_SSIM_heat_map(self):
        pass
    def FFID(self):
       
        pass
    def checkepoch(self,x,epoch,wd,k,f=False):
        if epoch % len(x) !=0:
            print("插值错误")
            return
       
        if len(x) < epoch:
            ind1 = np.arange(len(x))
            x = line_insert(ind1,x,int(epoch/len(x)))
        if f:
            return signal.savgol_filter(x,wd,k)
        else:
            return np.array(x)
    
def line_insert(x1,y1,k):
    lenE = len(y1) * k
    y2 = np.zeros(lenE+1)
    y2[x1*k]=y1
    y2[-1] = y1[-1]*2 - y1[-2]
    for i in range(lenE):
        if i % k == 0:
            continue
        p = int(i/k)*k
        y2[i] = (y2[p+k] - y2[p])*(i-p)/k + y2[p]
    return y2[:lenE]
    pass


